package com.raven.play.stream.test.javadsl.tutorial;

import akka.NotUsed;
import akka.stream.*;
import akka.stream.javadsl.*;
import com.raven.play.stream.test.BaseTest;
import org.junit.jupiter.api.Test;
import scala.collection.immutable.Seq;
import scala.jdk.javaapi.CollectionConverters;

import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CompletionStage;

public class $22CustomGraphShapes extends BaseTest {

    @Test
    public void test() {
        Source<Integer, NotUsed> source1 = Source.range(1, 10).throttle(1, Duration.of(1, ChronoUnit.SECONDS));
        Source<Integer, NotUsed> source2 = Source.range(20,30).throttle(3, Duration.of(1, ChronoUnit.SECONDS));

        Graph<Balance2x3, NotUsed> balance2x3 = GraphDSL.create(builder -> {
            UniformFanInShape<Integer, Integer> merge = builder.add(Merge.create(2));
            UniformFanOutShape<Integer, Integer> balance = builder.add(Balance.create(3));

            builder.from(merge).toFanOut(balance);

            return new Balance2x3(merge.in(0), merge.in(1),
                    balance.out(0), balance.out(1), balance.out(2));
        });

        RunnableGraph<NotUsed> balance2x3Graph = RunnableGraph.fromGraph(GraphDSL.create(builder -> {
            SourceShape<Integer> source1Shape = builder.add(source1);
            SourceShape<Integer> source2Shape = builder.add(source2);

            SinkShape<Object> sink1Shape = builder.add(createSink(1));
            SinkShape<Object> sink2Shape = builder.add(createSink(2));
            SinkShape<Object> sink3Shape = builder.add(createSink(3));

            Balance2x3 balance2x3Shape = builder.add(balance2x3);

            builder.from(source1Shape).toInlet(balance2x3Shape.in0);
            builder.from(source2Shape).toInlet(balance2x3Shape.in1);
            builder.from(balance2x3Shape.out0).to(sink1Shape);
            builder.from(balance2x3Shape.out1).to(sink2Shape);
            builder.from(balance2x3Shape.out2).to(sink3Shape);

            return ClosedShape.getInstance();
        }));

        balance2x3Graph.run(system);
    }

    private static Sink<Object, CompletionStage<Integer>> createSink(Integer index){
        return Sink.fold(0, (count, element) -> {
            log.info("{} received {}, count now is : {}", index, element, count);
            return count + 1;
        });
    }

    class Balance2x3 extends Shape {
        public final Inlet<Integer> in0;
        public final Inlet<Integer> in1;
        public final Outlet<Integer> out0;
        public final Outlet<Integer> out1;
        public final Outlet<Integer> out2;

        // private final Seq<Inlet<?>> inlets;
        // private final Seq<Outlet<?>> outlets;

        private final Seq<Inlet<?>> inlets;
        private final Seq<Outlet<?>> outlets;

        public Balance2x3() {
            this(Inlet.create("in0"), Inlet.create("in1"),
                    Outlet.create("out0"), Outlet.create("out1"), Outlet.create("out2"));
        }

        public Balance2x3(Inlet<Integer> in0, Inlet<Integer> in1,
                          Outlet<Integer> out0, Outlet<Integer> out1, Outlet<Integer> out2) {
            this.in0 = in0;
            this.in1 = in1;
            this.out0 = out0;
            this.out1 = out1;
            this.out2 = out2;

            // !!! Arrays.asList 和 CollectionConverters.asScala 不能合并到一行写, 否则会有类型推断错误
            List<Inlet<?>> inletList = Arrays.asList(in0, in1);
            List<Outlet<?>> outletList = Arrays.asList(out0, out1, out2);

            this.inlets = CollectionConverters.asScala(inletList).toSeq();
            this.outlets = CollectionConverters.asScala(outletList).toSeq();
        }

        @Override
        public Seq<Inlet<?>> inlets() {
            return inlets;
        }

        @Override
        public Seq<Outlet<?>> outlets() {
            return outlets;
        }

        @Override
        public Shape deepCopy() {
            return new Balance2x3(in0.carbonCopy(), in1.carbonCopy(),
                    out0.carbonCopy(), out1.carbonCopy(), out2.carbonCopy());
        }
    }
}
